import os
import shutil
import numpy as np
from PIL import Image
import h5py
import scipy
import matplotlib.pyplot as plt
import cv2

#加载图片
# image = cv2.imread(img_path)
# image = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
# image = Image.fromarray(image)
defect_tiff = np.array(Image.open(r'C:\Users\chanchan\Desktop\images\defect.tiff'))
defect_png = np.array(Image.open(r'C:\Users\chanchan\Desktop\images\defect.png'))
h5file = r'C:\Users\chanchan\Desktop\images\defect01.h5_origin'

def create_h5():
    #创建h5py文件
    h5file = r'C:\Users\chanchan\Desktop\images\defect_01.h5_origin'
    # if os.path.exists(h5file):
    #     shutil.rmtree(h5file)
    f = h5py.File(h5file, 'w')#生成H5文件

    #写如hpy文件
    #d1 = f.create_dataset("dset1", (20,), 'i')#deset1是数据集的name，(20,)代表数据集的shape，i代表的是数据集的元素类型
    f.create_dataset('tiff_image', data=defect_tiff)
    f.create_dataset('png_image', data=defect_png)
    f.create_dataset('defect_type', data=[89])
    f.close()

'''
import h5py
with h5py.File('file1.h5_origin','r') as f:
    for key in f.keys():
        print(f[key],key,f[key].name)
 
## f 表示h5文件的root目录，group是按字典的方式工作，f.keys()可以找到root下的所有group和dataset的key，然后通过key来访问
 
    data = f['X'][:]  # 读取表达量数据
    gene_names = f['var']['gene_name'][:]  # 读取基因名称
 
也可以对h5文件中的某个group进行
 
    dogs_group = f['dogs']
        for dkey in dogs_group.keys():
            print(dogs_group[dkey],dkey,dogs_group[dkey].name,dogs_group[dkey].value)
f.close()
 '''

def read_h5():
    #读取h5py
    h5file = r'C:\Users\chanchan\Desktop\images\defect_01.h5_origin'

    with h5py.File(h5file, 'r') as f:
        for key in f.keys():
            print(f[key], key, f[key].name)#f 表示h5文件的root目录，group是按字典的方式工作，f.keys()可以找到root下的所有group和dataset的key，然后通过key来访问

        tiff = f['tiff_image'][:]
        print(tiff, tiff.shape, type(tiff))#(680, 680) <class 'numpy.ndarray'>

        png = f['png_image'][:]
        print(png, png.shape, type(png)) #(720, 720) <class 'numpy.ndarray'>

        defect_type = f['defect_type'][0]
        print(defect_type, type(defect_type)) #89 <class 'numpy.int32'>

    f.close()

def test():
    tiff_path = r'C:\Users\chanchan\Desktop\images\defect.tiff'
    png_path = r'C:\Users\chanchan\Desktop\images\defect.png'
    imagepath = png_path
    img = None
    if '.tiff' in imagepath:
        img_cv = cv2.imread(imagepath, 0)
        img = Image.fromarray(img_cv)
    else:
        img = Image.open(imagepath)

    image = np.array(img)
    print(image, image.shape, type(image))#(680, 680) <class 'numpy.ndarray'>
    print(img, type(img))#<PIL.Image.Image image mode=L size=680x680 at 0x2D608D3BDD8> <class 'PIL.Image.Image'>


def build_dataset():
    image_dir = r'E:\proj\AI\Semi-Seg-Demo\data\myWood\image'
    mask_dir = r'E:\proj\AI\Semi-Seg-Demo\data\myWood\mask'
    h5_dir = r'../data/myWood/h5_origin'

    if os.path.exists(h5_dir):
        shutil.rmtree(h5_dir)
    os.mkdir(h5_dir)

    categories = os.listdir(image_dir)
    for category in categories:
        #遍历category下的image
        image_category_dir = os.path.join(image_dir, category)
        image_lists = os.listdir(image_category_dir)
        for image in image_lists:
            #获取文件名
            basename, _ = os.path.splitext(image)
            mask_image_name = basename + '_mask.png'

            #创建h5文件
            h5_file = os.path.join(h5_dir, basename + '.h5_origin')
            f = h5py.File(h5_file, 'w')  # 生成H5文件

            #将数据写入h5文件
            image_data = np.array(Image.open(os.path.join(image_dir, category, image)))
            mask_data = np.array(Image.open(os.path.join(mask_dir, category, mask_image_name)))
            f.create_dataset('image', data=image_data)
            f.create_dataset('mask', data=mask_data)
            f.create_dataset('category', data=[str(category)])
            f.close()


def read_dataset():
    h5_file = r'../data/myWood/h5_origin/color000.h5'
    with h5py.File(h5_file, 'r') as f:

        image = f['image'][:]
        print(image, image.shape, type(image))#(1024, 1024, 3) <class 'numpy.ndarray'>

        mask = f['mask'][:]
        print(mask, mask.shape, type(mask)) #(1024, 1024) <class 'numpy.ndarray'>

        category = f['category'][0]
        print(category, type(category)) #b'color' <class 'bytes'>

    f.close()






if __name__ == '__main__':
    # create_h5()
    # read_h5()
    # test()
    build_dataset()
    read_dataset()







